
# coding=utf-8
from __future__ import print_function, absolute_import, unicode_literals
from gm.api import *
import datetime
import pandas as pd
import numpy as np

def init(context):
    # 定义股票池数量
    context.num = 20
    # 波动率周期
    context.volatility = 252
    # 双周标记
    context.twoweekflg=True
    
    # 定时任务, 注册选股、卖出、买入三个定时任务
    schedule(schedule_func=algo, date_rule='1d', time_rule='09:15:00')
    # schedule(schedule_func=sell_algo, date_rule='1d', time_rule='09:30:00')
    # schedule(schedule_func=buy_algo, date_rule='1d', time_rule='09:31:30')

def algo(context):
    context.trading_status = False# 交易状态，TRUE时启动交易，FALSE关闭
    # 月份
    month = context.now.month # 提取当前月份数字
    # 过滤一月和四月（有持仓则平仓）
    if month==1 or month==4:
        order_close_all()
        return None
    # 上一个交易日，假如context.now是2024-05-06 09:15:00+08:00，那么last_date为2024-05-06
    last_date = get_previous_trading_date(exchange='SHSE', date=context.now)
    last_last_date = get_previous_trading_date(exchange='SHSE', date=last_date)
    # 那么next_date为 2024-05-08
    next_date = get_next_trading_date(exchange='SHSE', date=context.now)
    '''
    context.now 是一个表示当前日期的对象。
    context.now.weekday() 返回当前日期的星期几，
    其中星期一返回 0，星期二返回 1，依此类推，星期日返回 6。
    
    datetime.datetime.strptime(last_date, '%Y-%m-%d') 将字符串 last_date 解析为一个日期对象。
    .weekday() 返回该日期对象的星期几。
    '''
    # 如果当前日期的星期几小于或等于 last_date 的星期几，则执行下面的代码块。
    if context.now.weekday()<=datetime.datetime.strptime(last_date, '%Y-%m-%d').weekday():# 周一
        # context.twoweekflg 是一个布尔类型的标志位。在def init(context):中定义了
        # not context.twoweekflg 将 context.twoweekflg 的值取反
        # （即如果原来是 True，则变为 False；如果原来是 False，则变为 True）。
        context.twoweekflg=not context.twoweekflg
        # 如果 context.twoweekflg 为 True，则返回 None。
        if context.twoweekflg:
            return None
    '''
        上面这段代码的逻辑是：
        如果当前日期的星期几小于或等于 last_date 的星期几，则切换 context.twoweekflg 的值。
        如果切换后的 context.twoweekflg 为 True，则返回 None。
        这种逻辑通常用于某些周期性任务的控制，例如每两周执行一次某个操作。
    '''   
        # 获取A股代码（剔除停牌股、ST股、次新股（365天）,科创版）
    all_stock, all_stock_str = get_normal_stocks(context.now)
    all_stock_array=Arrsplits(all_stock,200)
    all_stock=[]
    # 加强：ROE>0
    '''=======================================================================================
    总结来说，这段代码的作用是：
    遍历 all_stock_array 中的每个子列表。
    对每个子列表中的股票代码查询其基本面数据（加权净资产收益率）。
    筛选出加权净资产收益率大于 0 的股票代码。
    将这些符合条件的股票代码添加到 all_stock 列表中。
    '''
    for arr in all_stock_array:
        fundamental=get_fundamentals(   # get_fundamentals 是一个函数，用于获取基本面数据。
            table='prim_finance_indicator', # 指定要查询的表。
            symbols=arr, # 指定要查询的股票代码列表。
            start_date=context.now, 
            end_date=context.now, 
            fields='ROEWEIGHTED', #  指定要查询的字段，ROEWEIGHTED 是加权净资产收益率（Return on Equity Weighted）的缩写。净资产收益率（ROE）
            limit=40000, #  指定要查询的字段，这里查询的是加权净资产收益率
            df=True #表示返回结果以 Pandas DataFrame 的形式返回。
        )
        # 筛选出加权净资产收益率大于 0 的股票。 [提取这些股票的股票代码。]
        # list(...) 将提取的股票代码转换为列表。
        stocks=list(fundamental[fundamental['ROEWEIGHTED']>0]['symbol']) 
        all_stock.extend(stocks) # 将筛选出的股票代码列表 stocks 添加到 all_stock 列表中。
        
    #=============================================================================================    
    # 获取所有股票市值,并按升序排序
    # 将 all_stock 列表分割成多个子列表，每个子列表最多包含 200 个元素，
    # 并将这些子列表存储在 all_stock_array 中。
    # 以便分批次处理数据，避免一次性处理大量数据导致的性能问题。
    all_stock_array=Arrsplits(all_stock,200)
    fundamental=None
    for arr in all_stock_array:
        if fundamental is None: #如果 fundamental 是 None，则对第一个子列表 arr 进行基本面数据查询，
                                #查询结果按总市值排序，并将结果赋值给 fundamental。
            fundamental = get_fundamentals(
                table='trading_derivative_indicator', #指定要查询的表。
                symbols=arr,    # 指定要查询的股票代码列表
                start_date=context.now, 
                end_date=context.now, 
                fields='TOTMKTCAP', # TOTMKTCAP 是总市值（Total Market Capitalization）的缩写。
                limit=40000, 
                df=True
            ).sort_values('TOTMKTCAP')  # 查询固定时间基本面数据
        else:# 如果 fundamental 不是 None，则对当前子列表 arr 进行基本面数据查询，查询结果按总市值排序，
             # 并将结果赋值给 fundamental_loc。
            fundamental_loc = get_fundamentals(
                table='trading_derivative_indicator', 
                symbols=arr, 
                start_date=context.now, 
                end_date=context.now, 
                fields='TOTMKTCAP', 
                limit=40000, 
                df=True
            ).sort_values('TOTMKTCAP')
            #  将 fundamental 和 fundamental_loc 合并。
            pd.concat([fundamental,fundamental_loc])
    # 获取前2N只股票,从 fundamental DataFrame 中选择前 context.num * 2 行的股票代码，并将这些股票代码存储在 to_buy 列表中。
    # [:context.num*2,:] 表示选择前 context.num * 2 行的所有列。[从选择的行中提取 symbol 列，即股票代码。]
    # list(...): 将提取的股票代码转换为列表。
    # to_buy 列表将包含 fundamental DataFrame 中前 context.num * 2 行的股票代码。
    to_buy = list(fundamental.iloc[:context.num*2,:]['symbol'])
    
    # # ================== 计算波动率===========================
    # 开始日期（计算日频收益率）
    start_date = get_previous_N_trading_date(context.now,counts=context.volatility+1)
    # 获取收盘价
    close = history_new(context, security=to_buy,frequency='1d',start_time=start_date,end_time=last_date,fields='eob,symbol,close',skip_suspended=True,fill_missing=None,adjust=ADJUST_PREV,adjust_end_time=last_date, df=True)
    # 日频收益率
    ret = (close/close.shift(1)-1).iloc[1:,:]
    # 波动率
    volatility = ret.std().sort_values()
    
    # 目标股票, [:context.num] 表示选择前 context.num 行的所有列。.index:从选择的行中提取索引（即行标签）。
    context.to_buy = list(volatility.iloc[:context.num].index)
    '''
            这行代码的作用是选择 volatility DataFrame 中前 context.num 行的索引，
            并将这些索引存储在 context.to_buy 列表中。
    '''
    
    print('{}:本次股票池有{}只:{}'.format(context.now,len(context.to_buy),context.to_buy))
    context.trading_status = True

# ============================================================================================
def get_normal_stocks(date,new_days=365):
    """
    获取目标日期date的A股代码（剔除停牌股、ST股、次新股（365天））
    :param date：目标日期
        date: 目标日期，可以是字符串格式
        （如 "2022-01-01" 或 "2022-01-01 12:00:00"）或 datetime 对象。
    :param new_days:新股上市天数，默认为365天
    """
    if isinstance(date,str) and len(date)==10: #isinstance(date, str) 检查 date 是否是字符串类型，如果是，则进一步检查其长度以确定日期的格式。
        date = datetime.datetime.strptime(date,"%Y-%m-%d")
    elif isinstance(date,str) and len(date)>10:
        date = datetime.datetime.strptime(date,"%Y-%m-%d %H:%M:%S")
    # 先剔除退市股、次新股和B股
    # 调用 get_instrumentinfos 函数，获取所有股票的代码、上市日期和退市日期，
    # 并将这些信息以Pandas DataFrame的形式返回，存储在变量 df_code 中。
    df_code = get_instrumentinfos(sec_types=SEC_TYPE_STOCK, fields='symbol, listed_date, delisted_date', df=True)
    
    
    #all_stocks = [code for code in df_code[(df_code['listed_date']<=date-datetime.timedelta(days=new_days))&(df_code['delisted_date']>date+datetime.timedelta(days=31))].symbol.to_list() if code[:6]!='SHSE.9' and code[:6]!='SZSE.2']
    all_stocks = [
    code        # 遍历前面筛选出的股票代码列表。
    for code in df_code[
        (df_code['listed_date'] <= date - datetime.timedelta(days=new_days)) & # 筛选出上市日期在 date 之前 new_days 天以上的股票。
        (df_code['delisted_date'] > date + datetime.timedelta(days=31)) # 筛选出退市日期在 date 之后 31 天以上的股票。
    ].symbol.to_list()  #将筛选后的 DataFrame 中的 symbol 列转换为列表。
    if code[:6] != 'SHSE.9' and code[:6] != 'SZSE.2' #进一步筛选出股票代码前 6 个字符不是 SHSE.9 和 SZSE.2 的股票。
]
    
    # 再剔除当前的停牌股和ST股
    history_ins = get_history_instruments(symbols=all_stocks, start_date=date, end_date=date, fields='symbol,sec_level, is_suspended', df=True)
    all_stocks = list(history_ins[(history_ins['sec_level']==1) & (history_ins['is_suspended']==0)]['symbol'])
    all_stocks = [code for code in all_stocks if code[:8]!='SHSE.688']
    # 将 all_stocks 列表中的所有元素用逗号 , 连接起来，形成一个字符串，并将这个字符串赋值给变量 all_stocks_str。
    all_stocks_str = ','.join(all_stocks)
    return all_stocks,all_stocks_str

# 分割数据的辅助函数
def Arrsplits(item, n):
    return [item[i:i+n] for i in range(0, len(item), n)]

# 获取指定日期 date 之前的第 N 个交易日。
# 具体来说，它会返回 date 之前的第 counts 个交易日，并且不包括 date 本身。
# date：目标日期，可以是 datetime 格式或其他可以转换为 pd.Timestamp 的格式。
# counts：历史回溯天数，默认为 1，即前一天。
# exchange：交易所代码，默认为 'SHSE'（上海证券交易所）。
def get_previous_N_trading_date(date,counts=1,exchange='SHSE'):
    """
    获取end_date前N个交易日,end_date为datetime格式，不包括date日期
    :param date：目标日期
    :param counts：历史回溯天数，默认为1，即前一天
    """
    date = pd.Timestamp(date) #将输入的 date 转换为 pd.Timestamp 格式，以便进行日期操作。
    previous_N_trading_date = get_trading_dates( # get_trading_dates获取从 start_date 到 date 之间的所有交易日列表。
        exchange=exchange, 
        # 计算起始日期 start_date，这是为了确保有足够的时间范围来找到前 counts 个交易日
        # max(counts + 30, counts * 3) 来确保起始日期足够远，以覆盖可能的非交易日（周末、节假日等）。
        start_date=date-datetime.timedelta(days=max(counts+30,counts*3)), 
        end_date=date
        )[-counts-1] # 从交易日列表中选择倒数第 counts + 1 个交易日，即 date 之前的第 counts 个交易日。
    return previous_N_trading_date # 返回计算得到的前 N 个交易日

def history_new(context,security,frequency,start_time,end_time,fields,skip_suspended=True,fill_missing=None,adjust=ADJUST_PREV, adjust_end_time='backtest_end_time', df=True):
    """
    分区间获取数据（以避免超出数据限制）(start_time和end_date为字符串,fields需包含eob和symbol,单字段)
    :param ：参数同history()参数一致，adjust_end_time默认为回测结束时间：context.backtest_end_time
    """
    if adjust_end_time=='backtest_end_time':
        adjust_end_time = context.backtest_end_time
    Data = pd.DataFrame()
    if frequency=='1d':
        trading_date = pd.Series(get_trading_dates(exchange='SZSE', start_date=start_time, end_date=end_time))
    else:
        trading_date = history('SHSE.000300', frequency=frequency, start_time=start_time, end_time=end_time, fields='eob', skip_suspended=skip_suspended, fill_missing=fill_missing, adjust=adjust, adjust_end_time=adjust_end_time, df=df)
        trading_date = trading_date['eob']
    space = 5
    if len(trading_date)<=space:
        Data = history(security, frequency=frequency, start_time=start_time, end_time=end_time, fields=fields, skip_suspended=skip_suspended, fill_missing=fill_missing, adjust=adjust, adjust_end_time=adjust_end_time, df=df)
    else:
        for n in range(int(np.ceil(len(trading_date)/space))):
            start = n*space
            end = start+space-1
            if end>=len(trading_date):
                data = history(security, frequency=frequency, start_time=trading_date.iloc[start], end_time=trading_date.iloc[-1], fields=fields, skip_suspended=skip_suspended, fill_missing=fill_missing, adjust=adjust, adjust_end_time=adjust_end_time, df=df)
            else:
                data = history(security, frequency=frequency, start_time=trading_date.iloc[start], end_time=trading_date.iloc[end], fields=fields, skip_suspended=skip_suspended, fill_missing=fill_missing, adjust=adjust, adjust_end_time=adjust_end_time, df=df)
            if len(data)>=33000:
                print('请检查返回数据量，可能超过系统限制，缺少数据！！！！！！！！！！')
            Data = pd.concat([Data,data])    
    Data.drop_duplicates(keep='first',inplace=True)
    if len(Data)>0:
        Data = Data.set_index(['eob','symbol'])
        Data = Data.unstack()
        Data.columns = Data.columns.droplevel(level=0)
    return Data

# =======================================================================================
if __name__ == '__main__':
    run(strategy_id='859c3b0a-1765-11ec-acc4-d0509949eed8',
    filename='main.py',
    mode=MODE_BACKTEST,
    # token='bf739ec28ad507f391a3c3e99ed0365007da88f8',
    token='ca9c4532786a122aa5ceb3fd726432862e046f91',
    backtest_start_time='2024-06-01 08:00:00',
    backtest_end_time='2024-06-20 15:00:00',
    backtest_adjust=ADJUST_PREV,
    backtest_initial_cash=10000000,
    backtest_commission_ratio=0.0001,
    backtest_slippage_ratio=0.0001)